#!/usr/bin/env python

from collections import defaultdict
import MySQLdb

def initialize_database (db_name,tbl_name):
  con = None
  db_server = 'localhost'
  db_username = 'root'
  db_password = 'mysqlroot'
  tbl_columns = 'Id INT PRIMARY KEY AUTO_INCREMENT, State VARCHAR(25), Population INT'
  create_db = "CREATE DATABASE " + db_name + ';'
  create_tbl = 'CREATE TABLE ' + tbl_name + '(' + tbl_columns + ')' + ';'

  try:
    con = MySQLdb.connection(host='localhost', user='root', passwd='mysqlroot');
    print create_db
    con.query(create_db)
    print 'USE Census;'
    con.query('USE Census;')
    print create_tbl
    con.query(create_tbl)

  except MySQLdb.Error, e:
    print "Error %d: %s" % (e.args[0],e.args[1])
    return False
  
  finally:    
    if con:    
      con.close()
      return True

def insert_into_database (db_name,tbl_name,state,population):
  con = None
  try:
    cmd = 'INSERT INTO ' + tbl_name + " (State,Population) Values ( '" + state + "', " + population +' ) ;'
    print cmd
    con = MySQLdb.connection(host='localhost', user='root', passwd='mysqlroot',db=db_name);
    con.query(cmd)
  except MySQLdb.Error, e:
    print "Error %d: %s" % (e.args[0],e.args[1])
    return False

  finally:    
    if con:    
      return True  

def query_total_population (db_name,tbl_name):
  con = None
  total_population = 0
  try:
    cmd = 'SELECT SUM(Population) as Total_Population FROM ' + tbl_name + " ;"
    print cmd
    con = MySQLdb.connection(host='localhost', user='root', passwd='mysqlroot',db=db_name);
    con.query(cmd)
    row = con.store_result()
    total_population = int(row.fetch_row()[0][0])

  except MySQLdb.Error, e:
    print "Error %d: %s" % (e.args[0],e.args[1])

  finally:    
    return total_population  

def print_state_population (db_name,tbl_name):
  con = None
  try:
    cmd = 'SELECT * FROM ' + tbl_name + " ;"
    print cmd
    print "Population by State from Database..."
    con = MySQLdb.connection(host='localhost', user='root', passwd='mysqlroot',db=db_name);
    con.query(cmd)
    result = con.store_result()
    for row in result.fetch_row(100):
      print row[1].ljust(25), ":", row[2]

  except MySQLdb.Error, e:
    print "Error %d: %s" % (e.args[0],e.args[1])

  finally:    
    return


def main():
  
  """
    [0,1]=State Code
    [2,3,4]=County Code
    [5]=Blank
    [6,7]=Age Group    
    [8:1000]=Population for 31 races and ethnicities
    [1002:]=County or county equivalent name
  """

  db_name = 'Census'
  tbl_name = 'Population'

  census_file = open('MR-CO.txt', 'r')
  state_population = defaultdict(int)
  for line in census_file:
    if line[2] == ' ':  #it's not county data
      state = line[1002:].strip()
      populationEntries = line[8:1000].split()
      population = sum([int(entry) for entry in populationEntries])
      state_population[state] = state_population[state] + population
  census_file.close()

  total = sum(state_population.values())

  print "Total Population: ",total
  print
  print "Population by State..."
  for state in sorted(state_population):
    print state.ljust(25),":",state_population[state]

  if initialize_database(db_name,tbl_name):
    for state in sorted(state_population):
      assert (insert_into_database(db_name,tbl_name,state,str(state_population[state]))) == True
  
  db_total = query_total_population(db_name,tbl_name)
  
  print 'Total Population from Database: ',db_total
  print_state_population(db_name,tbl_name)
  
  assert(total)==285230516
  assert(db_total)==285230516

if __name__ == '__main__':
    main()
